from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import AnyMessage,ToolMessage,SystemMessage,HumanMessage
from langgraph.graph import StateGraph


from pydantic import BaseModel,Field
from typing import TypedDict,Optional,Annotated

from sqlalchemy.orm import sessionmaker,declarative_base
from sqlalchemy import Column,String,Integer,create_engine

import json
import requests
import operator




class SearchQuery(BaseModel):
    query:str=Field(description="用于联网查询的问题")

@tool(args_schema=SearchQuery)
def fetch_real_time_info(query):
     """Get real-time Internet information"""
     url = "https://google.serper.dev/search"
     payload = json.dumps({
         "q": query,
         "num": 1,
     })
     headers = {
         'X-API-KEY': '22a84d67009121271e4a5eb21d809e11d3bc8d45',
         'Content-Type': 'application/json'
     }

     response = requests.post(url, headers=headers,payload=payload)
     # 将返回的JSON字符串转换为字典
     data=json.loads(response.text)
     if 'organic' in data:
         return json.dumps(data['organic'], ensure_ascii=False)  # 返回'organic'部分的JSON字符串
     else:
         return json.dumps({"error": "No organic results found"}, ensure_ascii=False)  # 如果没有'organic'键，返回错误信息

class WeatherLoca(BaseModel):
    location:str=Field(description="查询天气所在的城市")


@tool(args_schema=WeatherLoca)
def get_weather(location):
    """Call to get the current weather."""

    if location.lower() in ["北京"]:
        return "北京的温度是16度，天气晴朗。"
    elif location.lower() in ["重庆"]:
        return "重庆的温度是32度，天气晴朗。"
    else:
        return "不好意思，并未查询到具体的天气信息。"

class User(BaseModel):
    name:str=Field(description="用户名字")
    age:int=Field(description="用户年龄")
    phone:Optional[str]=Field(description="用户电话")
    email:Optional[str]=Field(description="用户邮箱")


Base = declarative_base()
class UserInfo(Base):
    __tablename__="users"
    id=Column(Integer,primary_key=True,autoincrement=True)
    name=Column(String(50))
    age=Column(Integer)
    phone=Column(String(20))
    email=Column(String(30))

# 数据库连接 URI，这里要替换成自己的Mysql 连接信息，以下是各个字段的对应解释：
# root：MySQL 数据库的用户名。
# 123456：MySQL 数据库的密码。
# localhost：MySQL 服务器的 IP 地址。
# langgraph：要连接的数据库的名称。
# charset=utf8mb4：设置数据库的字符集为 utf8mb4，支持更广泛的 Unicode 字符


DATABASE_URI = "mysql+pymysql://root:1qaz2wsx@localhost:3306/langgraph?charset=utf8mb4"
engine = create_engine(DATABASE_URI)
Session = sessionmaker(bind=engine)


@tool(args_schema=User)
def insert_db(name,age,phone,email):
    """Insert user information into the database, The required parameters are name, age, email, phone"""
    # 确保为每次操作创建新的会话
    session=Session()
    try:
        # 创建用户实例
        user_info=UserInfo(name=name,age=age,phone=phone,email=email)
        # 添加到会话
        session.add(user_info)
        # 提交事务
        session.commit()
        return {"messages":["成功添加到数据库"]}
    except Exception as e:
        session.rollback()
        return {"messages":[f"添加数据库失败，原因{e}"]}
    finally:
        session.close()


tools = [insert_db,get_weather,fetch_real_time_info]

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key = api_key,base_url=api_base,temperature=0)
llm = llm.bind_tools(tools)


def chat_with_model(state):
    print(state)
    messages = state["messages"]
    resp = llm.invoke(messages)
    return {"messages":[resp]}

class AgentState(TypedDict):
    messages:Annotated[list[AnyMessage],operator.add]

def execute_function(state:AgentState):
    print(state)
    tool_calls = state["messages"][-1].tool_calls
    results = []
    tools = [insert_db,get_weather,fetch_real_time_info]
    tools = {t.name:t for t in tools}
    for t in tool_calls:
        if not t["name"] in tools:
            result = "bad tool name, retry"
        else:
            result = tools[t["name"]].invoke(t["args"])
        results.append(ToolMessage(tool_call_id=t["id"],name=t["name"],content=str(result)))
    return {"messages":results}

def final_answer(state):
    """generate natural language responses"""
    messages = state["messages"][-1]
    return {"messages":[messages]}

# 请你基于现在得到的信息，进行总结，生成专业的回复，注意，请用中文回复
SYSTEM_PROMPT = """
Please summarize the information obtained so far and generate a professional response. Note, please reply in Chinese.
"""

def natural_response(state):
    """generate final language responses"""
    messages = state["messages"][-1]
    messages = [SystemMessage(content=SYSTEM_PROMPT)]+[HumanMessage(content=messages.content)]
    resp = llm.invoke(messages)
    return {"messages":[resp]}

def exists_function_calling(state:AgentState):
    result = state["messages"][-1]
    return len(result.tool_calls)>0

builder = StateGraph(AgentState)

builder.add_node("chat_with_model",chat_with_model)
builder.add_node("final_answer",final_answer)
builder.add_node("execute_function",execute_function)
builder.add_node("natural_response",natural_response)

builder.set_entry_point("chat_with_model")

builder.add_conditional_edges("chat_with_model",exists_function_calling,{True:"execute_function",False:"final_answer"})
builder.add_edge("execute_function","natural_response")
builder.add_edge("final_answer","natural_response")

graph = builder.compile()

# messages = [HumanMessage(content="请介绍一下你自己")]
# resp = graph.invoke({"messages":messages})
# print(resp)

query = "我叫王杏杏2，今年40岁，电话号码是175855245，邮箱是25458@qq.com"
input_message = {"messages":[HumanMessage(content=query)]}
result = graph.invoke(input_message)
print(f"result:{result}")